# Copyright (c) 2018 Presto Labs Pte. Ltd.
# Author: jaewon

import datetime
import sqlite3
import sys

import tabulate
from absl import app, flags


def to_date_str(date):
  if isinstance(date, (datetime.date, datetime.datetime)):
    return date.strftime('%Y%m%d')
  else:
    return date


class PortfolioDB(object):
  def __init__(self, db):
    self._conn = sqlite3.connect(db)
    self._cursor = self._conn.cursor()
    self._ensure_table()

  def __enter__(self):
    return self

  def __exit__(self, type, value, traceback):
    self.close()

  def __del__(self):
    self.close()

  def close(self):
    if self._cursor is not None:
      self._cursor.close()
      self._cursor = None
    if self._conn is not None:
      self._conn.close()
      self._conn = None

  def _ensure_table(self):
    self._cursor.execute('CREATE TABLE IF NOT EXISTS portfolio (\n'
                         '  id   INTEGER  PRIMARY KEY AUTOINCREMENT,\n'
                         '  date TEXT NOT NULL,\n'
                         '  name TEXT NOT NULL,\n'
                         '  UNIQUE(date, name)\n'
                         ');')
    self._conn.commit()

  def query(self, from_date_inclusive, to_date_inclusive=None):
    from_date = to_date_str(from_date_inclusive)
    to_date = to_date_str(to_date_inclusive or from_date_inclusive)

    self._cursor.execute("SELECT date, name FROM portfolio "
                         "WHERE '%s' <= date and date <= '%s';" % (from_date, to_date))
    kv_pairs = self._cursor.fetchall()

    res = {}
    for key, value in kv_pairs:
      assert from_date <= key <= to_date
      date = datetime.datetime.strptime(key, '%Y%m%d').date()
      if date not in res:
        res[date] = []
      res[date].append(value)

    return {k: sorted(v) for k, v in res.items()}

  def delete(self, date):
    date = to_date_str(date)
    self._cursor.execute("DELETE FROM portfolio WHERE date = '%s';" % date)
    self._conn.commit()

  def insert(self, date, name):
    date = to_date_str(date)
    self._cursor.execute("INSERT INTO portfolio (date, name) VALUES ('%s', '%s')" % (date, name))
    self._conn.commit()


class SimJobDB(object):
  def __init__(self, db):
    self._conn = sqlite3.connect(db)
    self._cursor = self._conn.cursor()
    self._ensure_table()

  def __enter__(self):
    return self

  def __exit__(self, type, value, traceback):
    self.close()

  def __del__(self):
    self.close()

  def close(self):
    if self._cursor is not None:
      self._cursor.close()
      self._cursor = None
    if self._conn is not None:
      self._conn.close()
      self._conn = None

  def _ensure_table(self):
    self._cursor.execute('CREATE TABLE IF NOT EXISTS sim_job_id (\n'
                         '  id        INTEGER  PRIMARY KEY AUTOINCREMENT,\n'
                         '  date      TEXT NOT NULL,\n'
                         '  job_id    TEXT NOT NULL,\n'
                         '  directory TEXT NOT NULL,\n'
                         '  UNIQUE(date, job_id)\n'
                         ');')
    self._conn.commit()

  def query(self, from_date_inclusive, to_date_inclusive=None):
    from_date = to_date_str(from_date_inclusive)
    to_date = to_date_str(to_date_inclusive or from_date_inclusive)

    self._cursor.execute("SELECT date, job_id, directory FROM sim_job_id "
                         "WHERE '%s' <= date and date <= '%s';" % (from_date, to_date))
    rows = self._cursor.fetchall()

    res = {}
    for row in rows:
      assert from_date <= row[0] <= to_date
      date = datetime.datetime.strptime(row[0], '%Y%m%d').date()
      if date not in res:
        res[date] = []
      res[date].append(row[1:])

    return {k: sorted(v) for k, v in res.items()}

  def insert(self, date, job_id, directory):
    date = to_date_str(date)
    self._cursor.execute(
        "INSERT INTO sim_job_id (date, job_id, directory) VALUES ('%s', '%s', '%s')" %
        (date, job_id, directory))
    self._conn.commit()


def main(argv):
  FLAGS = flags.FLAGS

  if len(argv) < 3:
    sys.stderr.write('Usage:   ops_util.py --db=<filename> <table> <from_date> <to_date>\n'
                     '* table: portflio or sim_job_id\n')
    return 1

  table = argv[1]

  if table == 'portfolio':
    db = PortfolioDB(FLAGS.db)
    from_date = argv[2]
    to_date = argv[3] if len(argv) >= 4 else argv[2]
    res = db.query(from_date, to_date)

    table = []
    for date in sorted(res.keys()):
      date_str = date.strftime('%Y%m%d')
      for idx, name in enumerate(res[date]):
        if idx == 0:
          table.append([date_str, name])
        else:
          table.append([None, name])
    print(tabulate.tabulate(table, headers=['Date', 'Portfolio']))

  elif table == 'sim_job_id':
    db = SimJobDB(FLAGS.db)
    from_date = argv[2]
    to_date = argv[3] if len(argv) >= 4 else argv[2]
    res = db.query(from_date, to_date)

    table = []
    for date in sorted(res.keys()):
      date_str = date.strftime('%Y%m%d')
      for job_id, directory in res[date]:
        table.append([date_str, job_id, directory])
    print(tabulate.tabulate(table, headers=['Date', 'Job ID', 'Directory']))

  else:
    sys.stderr.write('Unknown table: %s\n' % table)
    return 1


if __name__ == '__main__':
  flags.DEFINE_string('db', None, 'DB file')
  app.run(main)
